package com.qf.util;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.stereotype.Component;

import java.util.Collections;

/**
 * Bloom过滤器的工具类
 */
@Component
public class BloomUtil {

    @Autowired
    private StringRedisTemplate redisTemplate;

    //初始化布隆过滤器的脚本
    private String bloomInit = "local flag = redis.call('exists', KEYS[1]) if flag == 0 then return redis.call('bf.reserve', KEYS[1], ARGV[1], ARGV[2]) end";
    //添加数据到bloom过滤器
    private String bloomAdd = "return redis.call('bf.add', KEYS[1], ARGV[1])";
    //判断数据是否存在
    private String bloomExists = "return redis.call('bf.exists', KEYS[1], ARGV[1])";

    /**
     * 初始化布隆过滤器
     * @param key
     * @param v1
     * @param v2
     */
    public void initBloom(String key, double v1, Integer v2){
        redisTemplate.execute(new DefaultRedisScript(bloomInit),
                Collections.singletonList(key),
                v1 + "", v2 + "");
    }

    /**
     * 往布隆过滤器中添加元素
     * @param key
     * @param value
     * @return
     */
    public boolean addBloom(String key, String value){
        Long result = (Long) redisTemplate.execute(new DefaultRedisScript(bloomAdd, Long.class),
                Collections.singletonList(key),
                value);
        return result == 1;
    }

    /**
     * 判断元素是否存在于布隆过滤器
     * @param key
     * @param value
     * @return
     */
    public boolean isExists(String key, String value){
        Long result = (Long) redisTemplate.execute(new DefaultRedisScript(bloomExists, Long.class),
                Collections.singletonList(key),
                value);
        return result == 1;
    }
}
